import matplotlib.pyplot as plt
import numpy as np

# 读取数据
data = np.loadtxt('/home/sh/catkin_ws/src/ymbot_e_control/data/ymbot_e_ppo_data.txt', delimiter=',')

# 提取数据
timesteps = data[:, 0]
current_positions = data[:, 1:7]
target_positions = data[:, 7:13]
kps = data[:, 13:19]
kis = data[:, 19:25]
kds = data[:, 25:31]
torques = data[:, 31:37]

# 计算跟踪误差
errors = target_positions - current_positions

# 画出关节当前位置与目标位置跟踪效果图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(current_positions.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, current_positions[:, i], label=f'Joint {i+25} Current Position')
    ax.plot(timesteps, target_positions[:, i], '--', label=f'Joint {i+25} Target Position')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('Position')
    ax.set_title(f'Joint {i+25} Position Tracking')
    ax.legend()
fig.suptitle('Joint Position Tracking')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_position_tracking.png')

# 画出跟踪误差图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(errors.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, errors[:, i], label=f'Joint {i+25} Error')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('Error')
    ax.set_title(f'Joint {i+25} Tracking Error')
    ax.legend()
fig.suptitle('Joint Tracking Error')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_tracking_error.png')

# 画出输出力矩的图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(torques.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, torques[:, i], label=f'Joint {i+25} Torque')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('Torque')
    ax.set_title(f'Joint {i+25} Torques')
    ax.legend()
fig.suptitle('Joint Torques')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_torques.png')

# 画出每个关节的kp值图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(kps.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, kps[:, i], label=f'Joint {i+25} KP')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('KP')
    ax.set_title(f'Joint {i+25} KP Values')
    ax.legend()
fig.suptitle('Joint KP Values')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_kp_values.png')

# 画出每个关节的ki值图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(kis.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, kis[:, i], label=f'Joint {i+25} KI')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('KI')
    ax.set_title(f'Joint {i+25} KI Values')
    ax.legend()
fig.suptitle('Joint KI Values')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_ki_values.png')

# 画出每个关节的kd值图
fig, axs = plt.subplots(2, 3, figsize=(18, 10))
for i in range(kds.shape[1]):
    ax = axs[i // 3, i % 3]
    ax.plot(timesteps, kds[:, i], label=f'Joint {i+25} KD')
    ax.set_xlabel('Timestep')
    ax.set_ylabel('KD')
    ax.set_title(f'Joint {i+25} KD Values')
    ax.legend()
fig.suptitle('Joint KD Values')
plt.tight_layout()
# plt.savefig('/home/sh/catkin_ws/src/ymbot_e_control/plots/joint_kd_values.png')


plt.show()
